{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 10_Train_final_model_TransE_l2.ipynb\n",
    "#\n",
    "# created by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on March 11, 2023\n",
    "# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on March 11, 2023\n",
    "#\n",
    "# 该脚本展示了如何在 DRKG 上训练最终模型 (TransE_l2).\n",
    "#\n",
    "# 需要的包:\n",
    "#          torch\n",
    "#          dgl, version: 0.4.3\n",
    "#          dglke\n",
    "#\n",
    "# 需要的文件:\n",
    "#          ../../data/drkg/drkg.tsv\n",
    "#\n",
    "# 源教程链接: https://github.com/gnn4dr/DRKG/blob/master/embedding_analysis/Train_embeddings.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training DRKG Using TransE_l2\n",
    "\n",
    "这个 notebook 展示了如何在 DRKG 上训练最终模型 (TransE_l2)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们能使用 DGL-KE 命令训练 TransE_l2 模型, 关于如何使用 DGL-KE 的更多信息请参考 https://github.com/awslabs/dgl-ke.\n",
    "\n",
    "这里我们使用两个 GPU 训练模型."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- batch_size: **4096**\n",
    "\n",
    "- neg_sample_size: **256**\n",
    "\n",
    "- hidden_dim: 200, **400**\n",
    "\n",
    "- gamma: 6, **12**, 18\n",
    "\n",
    "- lr: 0.01, 0.05, **0.1**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading train triples....\n",
      "Finished. Read 5874261 train triples.\n",
      "|Train|: 5874261\n",
      "random partition 5874261 edges into 2 parts\n",
      "part 0 has 2937131 edges\n",
      "part 1 has 2937130 edges\n",
      "/home/luyanfeng/miniconda3/envs/drkg/lib/python3.8/site-packages/dgl/base.py:25: UserWarning: multigraph will be deprecated.DGL will treat all graphs as multigraph in the future.\n",
      "  warnings.warn(msg, warn_type)\n",
      "Total initialize time 14.657 seconds\n",
      "[proc 0][Train](20000/100000) average pos_loss: 0.33681009418407903\n",
      "[proc 1][Train](20000/100000) average pos_loss: 0.33603372640596446\n",
      "[proc 0][Train](20000/100000) average neg_loss: 0.4776868528932333\n",
      "[proc 1][Train](20000/100000) average neg_loss: 0.47734504776000974\n",
      "[proc 0][Train](20000/100000) average loss: 0.4072484735324979\n",
      "[proc 0][Train](20000/100000) average regularization: 0.06466307142076017\n",
      "[proc 1][Train](20000/100000) average loss: 0.40668938705176116\n",
      "[proc 0][Train] 20000 steps take 505.640 seconds\n",
      "[proc 0]sample: 69.508, forward: 233.611, backward: 67.174, update: 122.690\n",
      "[proc 1][Train](20000/100000) average regularization: 0.06463185143773931\n",
      "[proc 1][Train] 20000 steps take 505.643 seconds\n",
      "[proc 1]sample: 69.359, forward: 243.598, backward: 68.510, update: 121.206\n",
      "[proc 1][Train](40000/100000) average pos_loss: 0.30527374002337454\n",
      "[proc 0][Train](40000/100000) average pos_loss: 0.3052402902394533\n",
      "[proc 1][Train](40000/100000) average neg_loss: 0.44544545147418974\n",
      "[proc 1][Train](40000/100000) average loss: 0.3753595958575606\n",
      "[proc 0][Train](40000/100000) average neg_loss: 0.44521938632577657\n",
      "[proc 1][Train](40000/100000) average regularization: 0.06444449488669633\n",
      "[proc 0][Train](40000/100000) average loss: 0.3752298384025693\n",
      "[proc 1][Train] 20000 steps take 531.349 seconds\n",
      "[proc 1]sample: 72.332, forward: 257.262, backward: 70.913, update: 128.232\n",
      "[proc 0][Train](40000/100000) average regularization: 0.0644489256337285\n",
      "[proc 0][Train] 20000 steps take 531.365 seconds\n",
      "[proc 0]sample: 69.817, forward: 254.556, backward: 70.957, update: 129.061\n",
      "[proc 0][Train](60000/100000) average pos_loss: 0.3014769183412194\n",
      "[proc 1][Train](60000/100000) average pos_loss: 0.30142607345730066\n",
      "[proc 0][Train](60000/100000) average neg_loss: 0.440336926291883\n",
      "[proc 1][Train](60000/100000) average neg_loss: 0.44060979357808827\n",
      "[proc 0][Train](60000/100000) average loss: 0.37090692223310473\n",
      "[proc 0][Train](60000/100000) average regularization: 0.06389702703133225\n",
      "[proc 1][Train](60000/100000) average loss: 0.37101793343275785\n",
      "[proc 0][Train] 20000 steps take 567.014 seconds\n",
      "[proc 0]sample: 68.645, forward: 251.153, backward: 71.891, update: 167.651\n",
      "[proc 1][Train](60000/100000) average regularization: 0.06388911578804254\n",
      "[proc 1][Train] 20000 steps take 567.030 seconds\n",
      "[proc 1]sample: 72.331, forward: 260.131, backward: 73.181, update: 124.310\n",
      "[proc 0][Train](80000/100000) average pos_loss: 0.29979560211747885\n",
      "[proc 1][Train](80000/100000) average pos_loss: 0.29970968129634856\n",
      "[proc 0][Train](80000/100000) average neg_loss: 0.4381060902744532\n",
      "[proc 0][Train](80000/100000) average loss: 0.36895084620714186\n",
      "[proc 1][Train](80000/100000) average neg_loss: 0.43833865226060154\n",
      "[proc 0][Train](80000/100000) average regularization: 0.06355494891740382\n",
      "[proc 0][Train] 20000 steps take 624.333 seconds\n",
      "[proc 0]sample: 69.278, forward: 234.393, backward: 71.563, update: 248.791\n",
      "[proc 1][Train](80000/100000) average loss: 0.3690241667583585\n",
      "[proc 1][Train](80000/100000) average regularization: 0.06355245302133262\n",
      "[proc 1][Train] 20000 steps take 624.341 seconds\n",
      "[proc 1]sample: 69.220, forward: 238.671, backward: 68.679, update: 115.147\n",
      "[proc 1][Train](100000/100000) average pos_loss: 0.29883501267433166\n",
      "[proc 0][Train](100000/100000) average pos_loss: 0.29876760110557077\n",
      "[proc 1][Train](100000/100000) average neg_loss: 0.43715565409064294\n",
      "[proc 0][Train](100000/100000) average neg_loss: 0.43676474123448134\n",
      "[proc 1][Train](100000/100000) average loss: 0.3679953333452344\n",
      "[proc 1][Train](100000/100000) average regularization: 0.06331871912032366\n",
      "[proc 0][Train](100000/100000) average loss: 0.3677661712229252\n",
      "[proc 1][Train] 20000 steps take 514.099 seconds\n",
      "[proc 1]sample: 68.320, forward: 226.204, backward: 65.768, update: 122.233\n",
      "proc 1 takes 2742.461 seconds\n",
      "[proc 0][Train](100000/100000) average regularization: 0.06332669700756668\n",
      "[proc 0][Train] 20000 steps take 514.116 seconds\n",
      "[proc 0]sample: 68.527, forward: 223.580, backward: 66.386, update: 149.616\n",
      "proc 0 takes 2742.469 seconds\n",
      "training takes 2742.8752896785736 seconds\n",
      "Save model to ckpts/TransE_l2_All_DRKG_0\n"
     ]
    }
   ],
   "source": [
    "!DGLBACKEND=pytorch dglke_train --dataset All_DRKG --data_path ../../data/drkg \\\n",
    "--data_files drkg.tsv --format 'raw_udd_hrt' \\\n",
    "--model_name TransE_l2 \\\n",
    "--batch_size 4096 --neg_sample_size 256 --hidden_dim 400 \\\n",
    "--gamma 12.0 --lr 0.1 --max_step 100000 -adv --regularization_coef 1.00E-07 \\\n",
    "--gpu 0 1 --num_proc 2 --mix_cpu_gpu --async_update --force_sync_interval 1000 \\\n",
    "--log_interval 20000 --num_thread 32"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
